package com.example.da20;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * OpenAI聊天工具类，提供本地调用的多轮对话功能
 */
@Component
@Slf4j
public class OpenAiChat {

//    @Value("${openai.api.url:https://api.openai.com/v1/chat/completions}")
    private String apiUrl = "https://open.bigmodel.cn/api/paas/v4/chat/completions";
    
//    @Value("${openai.api.key:your-api-key}")
    private String apiKey = "3a9724814d0a4b8abcc15b4833d2b1e8.8P8Tmflr74oXnwdF";
    
//    @Value("${openai.api.model:gpt-3.5-turbo}")
    private String defaultModel;
    
    private final RestTemplate restTemplate = new RestTemplate();
    
    // 系统提示词
    private String systemPrompt;
    
    // 助手角色
    private String assistantRole;
    
    // 使用的模型
    private String model;
    
    // 对话历史
    private final List<Message> messageHistory = new ArrayList<>();
    
    /**
     * 默认构造函数
     */
    public OpenAiChat() {
        this.model = defaultModel;
    }
    
    /**
     * 带系统提示词和助手角色的构造函数
     * 
     * @param systemPrompt 系统提示词
     * @param assistantRole 助手角色
     */
    public OpenAiChat(String systemPrompt, String assistantRole) {
        this();
        this.setSystemPrompt(systemPrompt);
        this.setAssistantRole(assistantRole);
    }

    public void setApiUrl(String apiUrl){
        this.apiUrl = apiUrl;
    }
    
    /**
     * 带所有参数的构造函数
     * 
     * @param systemPrompt 系统提示词
     * @param assistantRole 助手角色
     * @param model 模型名称
     */
    public OpenAiChat(String systemPrompt, String assistantRole, String model) {
        this(systemPrompt, assistantRole);
        this.model = model;
    }
    
    /**
     * 设置系统提示词
     * 
     * @param systemPrompt 系统提示词
     * @return 当前实例，用于链式调用
     */
    public OpenAiChat setSystemPrompt(String systemPrompt) {
        this.systemPrompt = systemPrompt;
        
        // 清除旧的系统消息
        messageHistory.removeIf(message -> "system".equals(message.getRole()));
        
        // 添加新的系统消息
        if (systemPrompt != null && !systemPrompt.trim().isEmpty()) {
            messageHistory.add(new Message("system", systemPrompt));
        }
        
        return this;
    }
    
    /**
     * 设置助手角色
     * 
     * @param assistantRole 助手角色
     * @return 当前实例，用于链式调用
     */
    public OpenAiChat setAssistantRole(String assistantRole) {
        this.assistantRole = assistantRole;
        
        // 如果没有系统提示词，则将角色作为系统消息添加
        if ((systemPrompt == null || systemPrompt.trim().isEmpty()) 
                && assistantRole != null && !assistantRole.trim().isEmpty()) {
            
            // 移除旧的角色系统消息
            messageHistory.removeIf(message -> 
                "system".equals(message.getRole()) && 
                message.getContent() != null && 
                message.getContent().startsWith("You are a"));
            
            // 添加新的角色系统消息
            messageHistory.add(new Message("system", "You are a " + assistantRole + "."));
        }
        
        return this;
    }
    
    /**
     * 设置模型
     * 
     * @param model 模型名称
     * @return 当前实例，用于链式调用
     */
    public OpenAiChat setModel(String model) {
        this.model = model;
        return this;
    }
    
    /**
     * 添加用户消息并获取AI回复
     * 
     * @param userMessage 用户消息
     * @return AI回复内容
     */
    public String chat(String userMessage) {
        if (userMessage == null || userMessage.trim().isEmpty()) {
            throw new IllegalArgumentException("Message cannot be empty");
        }
        
        // 添加用户消息到历史
        messageHistory.add(new Message("user", userMessage));
        
        // 准备请求主体
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", model != null ? model : defaultModel);
        requestBody.put("messages", messageHistory);
        
        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", "Bearer " + apiKey);
        
        // 创建HTTP实体
        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);
        
        try {
            // 发送请求
            String response = restTemplate.postForObject(apiUrl, entity, String.class);
            JSONObject jsonResponse = JSON.parseObject(response);
            
            // 解析响应
            String assistantMessage = jsonResponse
                .getJSONArray("choices")
                .getJSONObject(0)
                .getJSONObject("message")
                .getString("content");
            
            // 添加助手回复到历史
            messageHistory.add(new Message("assistant", assistantMessage));
            
            return assistantMessage;
        } catch (Exception e) {
            log.error("Error calling OpenAI API: {}", e.getMessage());
            throw new RuntimeException("Failed to get response from OpenAI API: " + e.getMessage(), e);
        }
    }
    
    /**
     * 清除对话历史，重新开始对话
     * 保留系统提示词和角色设置
     */
    public void clearHistory() {
        String currentSystemPrompt = this.systemPrompt;
        String currentAssistantRole = this.assistantRole;
        
        // 清除所有历史
        messageHistory.clear();
        
        // 重新应用系统提示词和角色
        if (currentSystemPrompt != null && !currentSystemPrompt.trim().isEmpty()) {
            this.setSystemPrompt(currentSystemPrompt);
        } else if (currentAssistantRole != null && !currentAssistantRole.trim().isEmpty()) {
            this.setAssistantRole(currentAssistantRole);
        }
    }
    
    /**
     * 获取当前对话历史
     * 
     * @return 对话历史消息列表的副本
     */
    public List<Message> getMessageHistory() {
        return new ArrayList<>(messageHistory);
    }
}
